'''
Author: zhanwei xu
Date: 2023-12-22 15:50:28
LastEditors: zhanwei xu
LastEditTime: 2024-03-01 20:13:10
Description: 

Copyright (c) 2024 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''

from .utils import *
from fastapi.responses import StreamingResponse, JSONResponse
from datetime import datetime
import threading
import time
import os
import random

current_time = datetime.now()
print(current_time)

def save_count(chatbot):
    while True:
        current_time = datetime.now()
        if current_time.hour == 1 and current_time.minute == 0: #检查时间是否凌晨1点
            with open('count.txt', 'a') as f: # 打开文件以追加的方式保存count
                f.write(str(current_time) + ': ' + str(chatbot.count) + '\n')  # 保存count
            chatbot.count = 0  # 重置count
            time.sleep(60)  # 休眠60秒，以避免重复保存
        else:
            time.sleep(10) # 休眠10秒，然后再次检查时间

class Chatbot():  
    def __init__(self) -> None:
        self.model_name = "gpt-3.5-turbo"
        self.ip_traffic = {}
        self.ip_traffic_4 = {}
        self.api_key_index = 0
        self.api_keys = get_api_keys('api_keys.txt')
        self.api_key_index = random.randint(0,1)
        self.count = 0
        self.count_thread = threading.Thread(target=save_count, args=(self,))
        self.count_thread.daemon = True # 设置为守护线程，这样在主程序退出时，线程也会自动退出
        self.count_thread.start() # 启动线程

    async def chatbot(self, request):
        """
        This is an async function that serves as a chatbot and generates responses based on user input, with
        additional functionality for web search queries and saving conversation history.
        
        :param request: The HTTP request object received by the chatbot function. It contains information
        about the incoming request, such as the client's IP address and the message sent by the user
        :return: The `chatbot` function returns a `StreamingResponse` object with a JSON or text media type
        depending on the input message. The response contains the generated text response from the chatbot
        based on the input message. If certain conditions are met, such as the IP address exceeding the
        traffic limit or the message exceeding the maximum token limit, the function returns a different
        type of response, such as an error
        """
        
#        client_ip = request.client.host
        client_ip = request.headers.get("X-Real-IP")
        if True:
            if check_traffic(client_ip):
                print('asda')
                return StreamingResponse(generate_error_traffic(), media_type="text/plain")
            
        json_data = await request.json()
        # Extract message from request
        messages = json_data["message"]
        flag = 0
        if "filename" in messages[-1]:
            print('使用了图片')
            response = await async_request_gemini_vision(messages[-1])
        if 'assets' in messages[-1]:
            if len(messages[-1]['assets'])>0:
                print('使用了oss里的图片')
                print(len(messages[-1]['assets']))
                response = await async_request_gemini_vision(messages[-1])
            else:
                flag = 1
        else:
            flag =1

        if flag:
            content = []
            prefix_cotent = ''
            for message in messages:
                if message["role"]=="system":
                    content.append({
                        "role":"user",
                        "parts":[{"text":prefix_cotent+message["content"]}]
                    })
                    content.append({
                        "role":"model",
                        "parts":[{"text":"好的，我已了解了您的需求，请您告诉我具体的内容，我将按照您的要求进行回复"}]
                    })
                    prefix_cotent = message["content"]+"\n\n\n"
                else:
                    if message["role"]=="assistant":
                        role = "model"
                    else:
                        role= "user"
                    content.append({"role":role,
                    "parts":[{"text":prefix_cotent+message["content"]}]})
                    prefix_cotent= ''
            for index in range(len(content)):
                if content[index]["role"]=='user':
                    break
            content = content[index:]
            response = await async_request_gemini(content)
            
        # print(response.text)
        if response.status_code==200:
            return StreamingResponse(generate_text_gemini(response),media_type="text/event-stream" )
        else: 
            StreamingResponse(generate_error(),media_type="text/event-stream")

        return StreamingResponse(generate_text(response), media_type="text/event_stream" )


